#! /usr/bin/env python3

import os
import re
import time

import asyncio
import aiohttp
from faker import Faker
import requests
import aiofiles
import logging
from Crypto.Cipher import AES


fk = Faker(locale='zh-CN')
user_agent = fk.user_agent()
headers = {"User-Agent": user_agent}

logging.basicConfig(level=logging.INFO,
    format="%(asctime)s--%(levelname)s--%(message)s")

# 存放ts文件名
names = []

def download_m3u8(url):
    """下载m3u8文件"""
    os.system(f'curl -o index.m3u8 {url}')
    logging.info(f"{url} 下载完成")

async def download_ts(url, name: str, session, aes=None):
    ts_url = f"{url}{name}"
    async with aiofiles.open(f'my_video/{name}', 'wb') as f:
        async with session.get(ts_url, headers=headers) as response:
            content = await response.content.read()
            await f.write(aes.decrypt(content))
        logging.info(f"{name} 下载完成")

async def aio_download(url: str, m:int, n:int, aes=None):
    """从m3u8文件中获取ts文件地址并下载"""
    url = url.replace('index.m3u8', '')
    if not os.path.exists('my_video'):
        os.mkdir('my_video')
    tasks = []
    i = 0
    async with aiohttp.ClientSession() as session:
        for line in names[m:n]:
            line = line.strip() # 去除行尾的换行符
            i += 1
            names.append(line)
            task = asyncio.create_task(download_ts(url, line, session, aes))
            tasks.append(task)
        await asyncio.wait(tasks)

def merge(start: int, stop: int):
    os.chdir('my_video')
    s = ' '.join(names[start:stop])
    target = time.strftime('%m%d_%H%M%S', time.localtime(time.time()))
    os.system(f'cat {s} > {target}.mp4')  

def main():
    url = input("请输入m3u8的URI：")
    key = input("请输入key：")
    m = int(input(f"请输入下载的起点"))
    n = int(input(f"请输入下载的终点"))
    download_m3u8(url)
    # 删除注释行
    os.system('sed -i "/^#/d" index.m3u8')
    with open("index.m3u8", 'r', encoding="utf-8") as f:
        for line in f:
            names.append(line.strip())
    # 定义加密方法
    aes =AES.new(key=key, IV=b'0000000000000000', mode=AES.MODE_CBC)
    asyncio.run(aio_download(url, m, n, aes))
    merge(m, n)

if __name__ == "__main__":
    main()
